#pragma once

#include <torch/torch.h>
#include <nanoflann.hpp>

class CPointsTensor {
public:
    template <typename T>
    inline T* getIndex() {
        return m_KDTree != nullptr ? reinterpret_cast<T*>(m_KDTree) : buildIndex<T>();
    }

    template <typename T>
    inline T* buildIndex() {
        if (m_KDTree == nullptr) m_KDTree = static_cast<void*>(new T(3, *this, { 10 }));
        return reinterpret_cast<T*>(m_KDTree);
    }

    inline size_t kdtree_get_point_count() const { return m_Tensor.size(0); }
    inline float kdtree_get_pt(const size_t idx, const size_t dim) const {
        return m_Accessor[idx][dim];
    };
    template <class CBBOX>
    bool kdtree_get_bbox(CBBOX& /* bb */) const
    {
        return false;
    }

    template <typename T>
    void freeIndex() {
        if (m_KDTree != nullptr) {
            T* pTree = getIndex<T>();
            delete pTree;
            m_KDTree = nullptr;
        }
    }

    CPointsTensor(const torch::Tensor& tensor) : m_Tensor(tensor), m_Accessor(tensor.accessor<float, 2>()) {
    }
    ~CPointsTensor();

    torch::Tensor scales();

private:
    torch::Tensor m_Tensor;
    torch::TensorAccessor<float, 2> m_Accessor;

    void* m_KDTree = nullptr;

};

using KdTreeTensor = nanoflann::KDTreeSingleIndexAdaptor<
    nanoflann::L2_Simple_Adaptor<float, CPointsTensor>,
    CPointsTensor, 3, size_t
>;
